package com.xiashitech.agent.instrumentation.redis.jedis.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.constants.Instrumentation;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSJedisCallDepth;
import com.xiashitech.agent.instrumentation.redis.RedisData;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import static com.xiashitech.agent.instrumentation.redis.jedis.XSJedisSingletons.instrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSJedisInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("redis.clients.jedis.Jedis");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return named("redis.clients.jedis.Jedis");
    }

    @Override
    public void transform(TypeTransformer typeTransformer) {
        typeTransformer.applyAdviceToMethod(
                isMethod().and(named("hget"))
                        .and(takesArgument(0, isFinal().and(is(String.class))))
                        .and(takesArgument(1, isFinal().and(is(String.class))))
                        .and(takesArguments(2))
                        .and(returns(is(String.class)))
                        .and(isPublic()),
                XSJedisInstrumentation.class.getName() + "$HGetDataAdvice");
        typeTransformer.applyAdviceToMethod(
                isMethod().and(named("get"))
                        .and(takesArgument(0, is(String.class)))
                        .and(takesArguments(1))
                        .and(returns(is(String.class)))
                        .and(isPublic()),
                XSJedisInstrumentation.class.getName() + "$GetDataAdvice");
        typeTransformer.applyAdviceToMethod(
                isMethod().and(named("hset"))
                        .and(takesArgument(0, isFinal().and(is(String.class))))
                        .and(takesArgument(1, isFinal().and(is(String.class))))
                        .and(takesArgument(2, isFinal().and(is(String.class))))
                        .and(takesArguments(3))
                        .and(returns(named("java.lang.Long")))
                        .and(isPublic()),
                XSJedisInstrumentation.class.getName() + "$HSetDataAdvice");
        typeTransformer.applyAdviceToMethod(
                isMethod().and(named("set"))
                        .and(takesArgument(0, isFinal().and(is(String.class))))
                        .and(takesArgument(1, isFinal().and(is(String.class))))
                        .and(returns(named("java.lang.String")))
                        .and(isPublic()),
                XSJedisInstrumentation.class.getName() + "$SetDataAdvice");
        typeTransformer.applyAdviceToMethod(
                isMethod().and(named("setex"))
                        .and(takesArgument(0, is(String.class)))
                        .and(takesArgument(1, is(int.class)))
                        .and(takesArgument(2, is(String.class)))
                        .and(takesArguments(3))
                        .and(returns(named("java.lang.String")))
                        .and(isPublic()),
                XSJedisInstrumentation.class.getName() + "$PAndSetExDataAdvice");
        typeTransformer.applyAdviceToMethod(
                isMethod().and(named("psetex"))
                        .and(takesArgument(0, is(String.class)))
                        .and(takesArgument(1, is(long.class)))
                        .and(takesArgument(2, is(String.class)))
                        .and(takesArguments(3))
                        .and(returns(named("java.lang.String")))
                        .and(isPublic()),
                XSJedisInstrumentation.class.getName() + "$PAndSetExDataAdvice");
    }

    public static class HGetDataAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(0) String key,
                @Advice.Argument(1) String field,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSJedisCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            try {
                redisData = RedisData.builder().key(key == null ? null : key.getBytes()).field(field == null ? null : field.getBytes()).redisOpetion(Instrumentation.Redis_Opetion_Hget).build();
                Context parentContext = currentContext();
                if (redisData.getKey() == null || "".equals(redisData.getKey()) || redisData.getField() == null || "".equals(redisData.getField()) || !instrumenter().shouldStart(parentContext,redisData)) {
                    return;
                }
                parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                context = instrumenter().start(parentContext, redisData);
                context = BaggageUtil.syncBaggage().build().storeInContext(context);
                scope = context.makeCurrent();
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("HGetDataAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Return String value,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if (callDepth.decrementAndGet() > 0) {
                return;
            }
            if (scope == null) {
                return;
            }
            redisData.setValue(value == null ? null : value.getBytes());
            scope.close();
            instrumenter().end(context,redisData,redisData,null);
            durationPrinter.printElapseDuration("HGetDataAdvice e");
        }
    }

    public static class GetDataAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(0) String key,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSJedisCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            try {
                redisData = RedisData.builder().key(key == null ? null : key.getBytes()).redisOpetion(Instrumentation.Redis_Opetion_Get).build();
                Context parentContext = currentContext();
                if (redisData.getKey() == null || "".equals(redisData.getKey()) || !instrumenter().shouldStart(parentContext,redisData)) {
                    return;
                }
                parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                context = instrumenter().start(parentContext, redisData);
                context = BaggageUtil.syncBaggage().build().storeInContext(context);
                scope = context.makeCurrent();
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("GetDataAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Return String value,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if (callDepth.decrementAndGet() > 0) {
                return;
            }
            if (scope == null) {
                return;
            }
            redisData.setValue(value == null ? null : value.getBytes());
            scope.close();
            instrumenter().end(context,redisData,redisData,null);
            durationPrinter.printElapseDuration("GetDataAdvice e");
        }
    }

    public static class SetDataAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(0) String key,
                @Advice.Argument(1) String value,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSJedisCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            try {
                redisData = RedisData.builder().key(key == null ? null : key.getBytes()).value(value == null ? null : value.getBytes()).redisOpetion(Instrumentation.Redis_Opetion_Set).build();
                Context parentContext = currentContext();
                if (redisData.getKey() == null || "".equals(redisData.getKey()) || redisData.getValue() == null || "".equals(redisData.getValue()) || !instrumenter().shouldStart(parentContext,redisData)) {
                    return;
                }
                parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                context = instrumenter().start(parentContext, redisData);
                context = BaggageUtil.syncBaggage().build().storeInContext(context);
                scope = context.makeCurrent();
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("SetDataAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if (callDepth.decrementAndGet() > 0) {
                return;
            }
            if (scope == null) {
                return;
            }
            scope.close();
            instrumenter().end(context,redisData,redisData,null);
            durationPrinter.printElapseDuration("SetDataAdvice e");
        }
    }

    public static class HSetDataAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(0) String key,
                @Advice.Argument(1) String field,
                @Advice.Argument(2) String value,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSJedisCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            try {
                redisData = RedisData.builder().key(key == null ? null : key.getBytes()).field(field == null ? null : field.getBytes()).value(value == null ? null : value.getBytes()).redisOpetion(Instrumentation.Redis_Opetion_Hset).build();
                Context parentContext = currentContext();
                if (redisData.getKey() == null || "".equals(redisData.getKey()) || redisData.getField() == null || "".equals(redisData.getField()) || redisData.getValue() == null || "".equals(redisData.getValue()) || !instrumenter().shouldStart(parentContext,redisData)) {
                    return;
                }
                parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                context = instrumenter().start(parentContext, redisData);
                context = BaggageUtil.syncBaggage().build().storeInContext(context);
                scope = context.makeCurrent();
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("HSetDataAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if (callDepth.decrementAndGet() > 0) {
                return;
            }
            if (scope == null) {
                return;
            }
            scope.close();
            instrumenter().end(context,redisData,redisData,null);
            durationPrinter.printElapseDuration("HSetDataAdvice e");
        }
    }

    public static class PAndSetExDataAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(0) String key,
                @Advice.Argument(1) long time,
                @Advice.Argument(2) String value,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSJedisCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            try {
                redisData = RedisData.builder().key(key == null ? null : key.getBytes()).time(time).value(value.getBytes()).redisOpetion(Instrumentation.Redis_Opetion_PAndSetEx).build();
                Context parentContext = currentContext();
                if (redisData.getKey() == null || "".equals(redisData.getKey()) || redisData.getValue() == null || "".equals(redisData.getValue()) || !instrumenter().shouldStart(parentContext,redisData)) {
                    return;
                }
                parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                context = instrumenter().start(parentContext, redisData);
                context = BaggageUtil.syncBaggage().build().storeInContext(context);
                scope = context.makeCurrent();
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("PAndSetExDataAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("redisData") RedisData redisData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSJedisInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if (callDepth.decrementAndGet() > 0) {
                return;
            }
            if (scope == null) {
                return;
            }
            scope.close();
            instrumenter().end(context,redisData,redisData,null);
            durationPrinter.printElapseDuration("PAndSetExDataAdvice e");
        }
    }
}
